#include <inference_engine.hpp> 
#include <iostream>
#include <string> 
#include <vector>
#include <opencv2/opencv.hpp>
#include <fstream>
using namespace InferenceEngine;
using namespace std;

string inputName;
string outputName;
InferRequest inferReq;
vector<wstring> labels;
//初试化模型相关参数
void initModel(string xml,string bin,string plugin="plugins.xml") { 
	try {
		Core ie(plugin);
		CNNNetReader network_reader;
		network_reader.ReadNetwork(xml);
		network_reader.ReadWeights(bin);
		network_reader.getNetwork().setBatchSize(1);
		CNNNetwork network = network_reader.getNetwork();
		InputInfo::Ptr input_info = network.getInputsInfo().begin()->second;
		inputName = network.getInputsInfo().begin()->first;

		input_info->getPreProcess().setResizeAlgorithm(RESIZE_BILINEAR);
		input_info->setLayout(Layout::NCHW);
		input_info->setPrecision(Precision::U8);

		DataPtr output_info = network.getOutputsInfo().begin()->second;
		outputName = network.getOutputsInfo().begin()->first;

		output_info->setPrecision(Precision::FP32);

		ExecutableNetwork executable_network = ie.LoadNetwork(network, "CPU");

		inferReq = executable_network.CreateInferRequest();
	}catch (const std::exception & ex) {
		std::cerr << ex.what() << std::endl; 
	}
}
//Mat 转Blob
void  matU8ToBlob(const cv::Mat& orig_image, InferenceEngine::Blob::Ptr& blob, int batchIndex=0) {
	InferenceEngine::SizeVector blobSize = blob->getTensorDesc().getDims();
	const size_t width = blobSize[3];
	const size_t height = blobSize[2];
	const size_t channels = blobSize[1];
	uint8_t* blob_data = blob->buffer().as<uint8_t*>();

	cv::Mat resized_image(orig_image);
	if (static_cast<int>(width) != orig_image.size().width ||
		static_cast<int>(height) != orig_image.size().height) {
		cv::resize(orig_image, resized_image, cv::Size(width, height));
	}

	int batchOffset = batchIndex * width * height * channels;

	for (size_t c = 0; c < channels; c++) {
		for (size_t h = 0; h < height; h++) {
			for (size_t w = 0; w < width; w++) {
				blob_data[batchOffset + c * width * height + h * width + w] =
					resized_image.at<cv::Vec3b>(h, w)[c];
			}
		}
	}
}
//读取label
void readLabel(string labelPath)
{
	std::wstring_convert<std::codecvt_utf8<wchar_t>> conv;
	ifstream in(labelPath.c_str());
	string line;
	if (in) { // 有该文件 
		while (getline(in, line)) { // line中不包括每行的换行符 
			wstring wb = conv.from_bytes(line);
			labels.push_back(wb);
		} 
	} else { // 没有该文件 
		cout << "no such file:" << labelPath << endl;
	} 
}
//前向计算
wstring infer(cv::Mat rgb,float& rtP) {
	Blob::Ptr imgBlob = inferReq.GetBlob(inputName);
	matU8ToBlob(rgb, imgBlob);
	inferReq.Infer();
	Blob::Ptr output = inferReq.GetBlob(outputName);
	float* logits = output->buffer().as<InferenceEngine::PrecisionTrait<InferenceEngine::Precision::FP32>::value_type*>();

	int maxIdx = 0;
	float maxP = 0;
	int nclasses = labels.size();//1001类
	float sum = 1;
	//softmax
	for (int i = 0; i < nclasses; i++) {
		logits[i] = exp(logits[i]);
		sum = sum + logits[i];
		if (logits[i] > maxP) {
			maxP = logits[i];
			maxIdx = i;
		}
	}
	
	rtP = maxP / sum; 
	return labels[maxIdx];
}
//测试
int main()
{
	string xml = "../model/mobilenet_v1_1.0_224_frozen.xml";
	string bin = "../model/mobilenet_v1_1.0_224_frozen.bin";
	string plugin = "../model/plugins.xml";
	string label = "../model/labels.txt";
	string testImg = "../model/test.png";
	initModel(xml, bin, plugin);
	readLabel(label);

	cv::Mat test = cv::imread(testImg);
	cv::Mat rgb;
	cv::cvtColor(test,rgb, cv::COLOR_BGR2RGB);
	float p;
	wstring cls = infer(rgb, p);
	std::wcout.imbue(std::locale("chs"));
	wcout << "类别：" << cls << ",概率：" << p << endl;
}
 
